import tushare as ts
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import mplfinance as mpf
# 初始化 TuShare Pro 接口
pro = ts.pro_api('31d3906685a1b0d31d52b4478cc1448a5ee4235abe9f55ab0a8943b9')

# 查询单支股票的日线数据
df_single = pro.daily(ts_code="000001.SZ", start_date='20220101', end_date='20250421', fields='ts_code,trade_date,open,high,low,close,vol')
df_single['trade_date'] = pd.to_datetime(df_single['trade_date'])
df_single.set_index('trade_date', inplace=True)

# 计算移动平均线
df_single['MA5'] = df_single['close'].rolling(window=5).mean()
df_single['MA20'] = df_single['close'].rolling(window=20).mean()

# 生成交易信号
df_single['signal'] = 0
df_single.loc[df_single['MA5'] > df_single['MA20'], 'signal'] = 1  # 金叉，买入信号
df_single.loc[df_single['MA5'] < df_single['MA20'], 'signal'] = -1  # 死叉，卖出信号

# 去除信号为0的行（无交易信号）
df_single = df_single[df_single['signal'] != 0].copy()

# 打印交易信号
print(df_single[['close', 'MA5', 'MA20', 'signal']])
# 查询多支股票的日线数据
stocks = ['000001.SZ', '000002.SZ', '000003.SZ', '000004.SZ']
df = pro.daily(ts_code=','.join(stocks), start_date='20220101', end_date='20250421', fields='ts_code,trade_date,close')
df['trade_date'] = pd.to_datetime(df['trade_date'])
df.set_index('trade_date', inplace=True)
df_wide = df.pivot(columns='ts_code', values='close')
df_wide.sort_index(inplace=True)

# 计算每支股票的日收益率和累计收益率
returns = df_wide.pct_change()
cumulative_returns = (1 + returns).cumprod() - 1
total_returns = cumulative_returns.iloc[-1]

# 将股票按总收益率排序
sorted_stocks = sorted(total_returns.items(), key=lambda x: x[1], reverse=True)

print("按总收益率排序的股票：")
for stock, return_rate in sorted_stocks:
    print(f"{stock}: {return_rate:.2%}")
# 设置止损和止盈百分比
stop_loss_pct = -0.05  # 止损5%
take_profit_pct = 0.10  # 止盈10%

# 假设我们已经买入股票，记录买入价格
buy_price = df_single.loc[df_single['signal'] == 1, 'close'].iloc[0]

# 计算止损和止盈价格
stop_loss_price = buy_price * (1 + stop_loss_pct)
take_profit_price = buy_price * (1 + take_profit_pct)

print(f"买入价格：{buy_price:.2f}")
print(f"止损价格：{stop_loss_price:.2f}")
print(f"止盈价格：{take_profit_price:.2f}")
# 初始化资金和持仓
initial_funds = 100000  # 初始资金10万元
funds = initial_funds
positions = 0  # 持仓数量

# 回测
for index, row in df_single.iterrows():
    if row['signal'] == 1 and funds > 0:  # 买入信号
        positions = funds / row['close']  # 全仓买入
        funds = 0
    elif row['signal'] == -1 and positions > 0:  # 卖出信号
        funds = positions * row['close']  # 卖出持仓
        positions = 0

# 计算最终资产
final_assets = funds + positions * df_single['close'].iloc[-1]
total_return = (final_assets - initial_funds) / initial_funds

print(f"初始资金：{initial_funds:.2f}")
print(f"最终资产：{final_assets:.2f}")
print(f"总收益率：{total_return:.2%}")
# 绘制K线图和交易信号
df_single.reset_index(inplace=True)
df_single.rename(columns={'open': 'Open', 'high': 'High', 'low': 'Low', 'close': 'Close', 'vol': 'Volume'}, inplace=True)
df_single.set_index('trade_date', inplace=True)  # 确保索引是DatetimeIndex

# 创建买入和卖出信号的列，填充NaN
df_single['buy_signals'] = np.nan
df_single['sell_signals'] = np.nan

# 填充买入和卖出信号
df_single.loc[df_single['signal'] == 1, 'buy_signals'] = df_single['Close']
df_single.loc[df_single['signal'] == -1, 'sell_signals'] = df_single['Close']

# 绘制K线图
apds = [
    mpf.make_addplot(df_single['MA5'], color='fuchsia', title='MA5'),
    mpf.make_addplot(df_single['MA20'], color='b', title='MA20'),
    mpf.make_addplot(df_single['buy_signals'], scatter=True, markersize=100, marker='^', color='g', title='Buy Signals'),
    mpf.make_addplot(df_single['sell_signals'], scatter=True, markersize=100, marker='v', color='r', title='Sell Signals')
]

mpf.plot(df_single, type='candle', addplot=apds, title='000001.SZ Price and Signals')

# 绘制收益率曲线
plt.figure(figsize=(12, 6))
plt.plot(df_single.index, df_single['Close'].pct_change().cumsum(), label='Cumulative Return')
plt.xlabel('Date')
plt.ylabel('Cumulative Return')
plt.title('Cumulative Return Curve')
plt.legend()
plt.show()
# 可视化收益率曲线
plt.figure(figsize=(14, 7))
for column in cumulative_returns.columns:
    plt.plot(cumulative_returns.index, cumulative_returns[column], label=column)

plt.title('Cumulative Returns of Multiple Stocks')
plt.xlabel('Date')
plt.ylabel('Cumulative Return')
plt.legend()
plt.grid(True)
plt.show()